# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jaewon

from experimental.prophet.graph import graph
from experimental.prophet.ops.concat import concat
from experimental.prophet.ops.constant import constant
from experimental.prophet.ops import elemwise_math
from experimental.prophet.ops.util import to_tuple


def linear_regression_predict(vars, coef_dict):
  vars = {v.name: v.cast(float) for v in to_tuple(vars)}
  if 'const' in coef_dict and 'const' not in vars:
    vars['const'] = constant(1., dtype=float)

  assert sorted(vars.keys()) == sorted(coef_dict.keys())

  x_vars = []
  coef = []
  for k in coef_dict.keys():
    x_vars.append(vars[k])
    coef.append(coef_dict[k])

  x = concat(*x_vars, dtype=float)
  coef = constant(coef, dtype=float)
  return elemwise_math.sum(x * coef)
